#! /usr/bin/env python3
# coding: utf-8
"""
References:
[1]Fast Multichannel Nonnegative Matrix Factorization with Directivity-Aware Jointly-Diagonalizable
Spatial Covariance Matrices for Blind Source Separation
[2]Fast Multichannel Source Separation Based on Jointly Diagonalizable Spatial Covariance Matrices
"""

import numpy as np
import chainer
import sys, os
import librosa
import soundfile as sf
import time
import pickle as pic

from configure_FastModel import *
from FastFCA import FastFCA
from FastMNMF import FastMNMF

try:
    from chainer import cuda
    FLAG_GPU_Available = True
except:
    print("---Warning--- You cannot use GPU acceleration because chainer or cupy is not installed")


class FastMNMF2(FastMNMF):  # 继承父类
    """ Blind Source Separation Using Fast Multichannel Nonnegative Matrix Factorization (FastMNMF)

    X_FTM: the observed complex spectrogram
    Q_FMM: diagonalizer that converts a spatial covariance matrix (SCM) to a diagonal matrix
    G_NFM: diagonal elements of the diagonalized SCMs
    W_NFK: basis vectors for each source
    H_NKT: activations for each source
    lambda_NFT: power spectral densities of each source (W_NFK @ H_NKT)
    Qx_power_FTM: power spectra of Qx
    Y_FTM: \sum_n lambda_NFT G_NM
    """

    def __init__(self, n_source=2, n_basis=8, xp=np, init_SCM="unit", dl_ratio=0):
        """ initialize FastMNMF

        Parameters:
        -----------
            n_source: int
                the number of sources
            n_iteration: int
                the number of iteration to update all variables
            n_basis: int
                the number of bases of each source
            init_SCM: str
                how to initialize covariance matrix {unit, obs, ILRMA}
        """
        super(FastMNMF2, self).__init__(n_source=n_source, xp=xp, init_SCM=init_SCM, dl_ratio=dl_ratio)
        self.n_basis = n_basis
        self.method_name = "FastMNMF2"


    def set_parameter(self, n_source=None, n_basis=None, init_SCM=None):
        """ set parameters

        Parameters:
        -----------
            n_source: int
            n_iteration: int
            n_basis: int
            init_SCM: str
                how to initialize covariance matrix {unit, obs, ILRMA}
        """
        super(FastMNMF2, self).set_parameter(n_source=n_source, init_SCM=init_SCM)  # super()调用父类
        if n_basis != None:
            self.n_basis = n_basis


    def update(self):
        self.update_WH()
        self.update_CovarianceDiagElement()
        self.udpate_Diagonalizer()
        self.normalize()


    def update_WH(self):
        tmp1_NFT = (self.G_NM[:, None, None] * (self.Qx_power_FTM / (self.Y_FTM ** 2))[None]).sum(axis=3)  # numerator of (34)(35) in [1]
        tmp2_NFT = (self.G_NM[:, None, None] / self.Y_FTM[None]).sum(axis=3)  # denominator of (34)(35) in [1]
        a_W = (self.H_NKT[:, None] * tmp1_NFT[:, :, None]).sum(axis=3)  # N F K T -> N F K
        b_W = (self.H_NKT[:, None] * tmp2_NFT[:, :, None]).sum(axis=3)  # N F K T -> N F K
        a_H = (self.W_NFK[..., None] * tmp1_NFT[:, :, None]).sum(axis=1) # N F K T -> N K T
        b_H = (self.W_NFK[..., None] * tmp2_NFT[:, :, None]).sum(axis=1) # N F K T -> N K T
        self.W_NFK = self.W_NFK * self.xp.sqrt(a_W / b_W)  # (34) of [1]
        self.H_NKT = self.H_NKT * self.xp.sqrt(a_H / b_H)  # (35) of [1]

        self.lambda_NFT = self.W_NFK @ self.H_NKT + EPS
        self.Y_FTM = (self.lambda_NFT[..., None] * self.G_NM[:, None, None]).sum(axis=0)


    def update_CovarianceDiagElement(self):
        a_1 = (self.lambda_NFT[..., None] * (self.Qx_power_FTM / (self.Y_FTM ** 2))[None]).sum(axis=2).sum(axis=1)  # N F T M
        b_1 = (self.lambda_NFT[..., None] / self.Y_FTM[None]).sum(axis=2).sum(axis=1)
        self.G_NM = self.G_NM * self.xp.sqrt(a_1 / b_1)  # (36) of [1]
        self.G_NM += EPS
        self.Y_FTM = (self.lambda_NFT[..., None] * self.G_NM[:, None, None]).sum(axis=0)


    def normalize(self):
        phi_F = self.xp.trace(self.Q_FMM * self.Q_FMM.conj(), axis1=1, axis2=2).real / self.n_mic  # (37) of [1]
        self.Q_FMM = self.Q_FMM / self.xp.sqrt(phi_F)[:, None, None]
        self.W_NFK = self.W_NFK / phi_F[None, :, None]

        mu_N = (self.G_NM).sum(axis=1).real  # (38) of [1]
        self.G_NM = self.G_NM / mu_N[:, None]
        self.W_NFK = self.W_NFK * mu_N[:, None, None]

        nu_NK = self.W_NFK.sum(axis=1)  # (28) of [1]
        self.W_NFK = self.W_NFK / nu_NK[:, None]
        self.H_NKT = self.H_NKT * nu_NK[:, :, None]
        self.lambda_NFT = self.W_NFK @ self.H_NKT + EPS

        self.reset_variable()


    def reset_variable(self):
        if self.xp == np:  # below(31) of [1]
            self.Qx_power_FTM = self.xp.abs((self.Q_FMM[:, None] @ self.X_FTM[:, :, :, None])[:, :, :, 0]) ** 2
        else:
            self.Qx_power_FTM = self.xp.abs((self.Q_FMM[:, None] * self.X_FTM[:, :, None]).sum(axis=3)) ** 2
        self.Y_FTM = (self.lambda_NFT[..., None] * self.G_NM[:, None, None]).sum(axis=0)


    def initialize_covarianceMatrix(self):
        if "unit" in self.init_SCM:
            # tile相当于MATLAB的repmat Q_FMM是diagonalizer
            self.Q_FMM = self.xp.tile(self.xp.eye(self.n_mic), [self.n_freq, 1, 1]).astype(self.xp.complex)
            self.G_NM = self.xp.ones([self.n_source, self.n_mic], dtype=self.xp.float) * 1e-2
            for m in range(self.n_mic):
                self.G_NM[m % self.n_source, m] = 1  # diagonal initialization
        elif "obs" in self.init_SCM:
            mixture_covarianceMatrix_FMM = self.XX_FTMM.sum(axis=1) / (self.xp.trace(self.XX_FTMM, axis1=2, axis2=3).sum(axis=1))[:, None, None]
            eig_val, eig_vec = np.linalg.eigh(self.convert_to_NumpyArray(mixture_covarianceMatrix_FMM))  # 特征值特征向量
            self.Q_FMM = self.xp.asarray(eig_vec).transpose(0, 2, 1).conj()
            self.G_NM = self.xp.ones([self.n_source, self.n_mic], dtype=self.xp.float) / self.n_mic
            self.G_NM[0] = self.xp.asarray(eig_val[0])
        elif "ILRMA" in self.init_SCM:
            sys.path.append("../Rank1_Model")
            from ILRMA import ILRMA
            ilrma = ILRMA(n_basis=2, init_SCM="unit", xp=self.xp)
            ilrma.load_spectrogram(self.X_FTM)
            ilrma.solve(n_iteration=15, save_likelihood=False, save_wav=False, save_path="./", interval_save_parameter=1000)
            separated_spec_power = self.xp.abs(ilrma.separated_spec).mean(axis=(1, 2))
            self.Q_FMM = ilrma.SeparationMatrix_FMM
            self.G_NM = self.xp.ones([self.n_source, self.n_mic], dtype=self.xp.float) * 1e-2
            for n in range(self.n_source):
                self.G_NM[n, separated_spec_power.argmax()] = 1
                separated_spec_power[separated_spec_power.argmax()] = 0
        else:
            print("Please specify how to initialize covariance matrix {unit, obs}")
            raise ValueError

        self.normalize()


    def separate_FastWienerFilter(self, source_index=None, mic_index=MIC_INDEX):  # 要改
        Qx_FTM = (self.Q_FMM[:, None] * self.X_FTM[:, :, None]).sum(axis=3)
        if source_index != None:
            diagonalizer_inv_FMM = self.calculateInverseMatrix(self.Q_FMM)
            self.separated_spec = self.convert_to_NumpyArray((diagonalizer_inv_FMM[:, None] @ (Qx_FTM * ( (self.lambda_NFT[source_index, :, :, None] * self.G_NM[source_index, None, None]) / (self.lambda_NFT[..., None]* self.G_NM[:, None, None]).sum(axis=0) ) )[..., None])[:, :, mic_index, 0])
        else:
            for n in range(self.n_source):
                diagonalizer_inv_FMM = self.calculateInverseMatrix(self.Q_FMM)
                tmp = self.convert_to_NumpyArray((diagonalizer_inv_FMM[:, None] @ (Qx_FTM * ( (self.lambda_NFT[n, :, :, None] * self.G_NM[n, None, None]) / (self.lambda_NFT[..., None]* self.G_NM[:, None, None]).sum(axis=0) ) )[..., None])[:, :, mic_index, 0])
                if n == 0:
                    self.separated_spec = np.zeros([self.n_source, tmp.shape[0], tmp.shape[1]], dtype=np.complex)
                self.separated_spec[n] = tmp


    def save_parameter(self, fileName):
        param_list = [self.lambda_NFT, self.G_NM, self.Q_FMM]
        if self.xp != np:
            param_list = [self.convert_to_NumpyArray(param) for param in param_list]
        pic.dump(param_list, open(fileName, "wb"))


    def load_parameter(self, fileName):
        param_list = pic.load(open(fileName, "rb"))
        if self.xp != np:
            param_list = [cuda.to_gpu(param) for param in param_list]
        self.lambda_NFT, self.G_NM, self.Q_FMM = param_list

        self.n_source, self.n_freq, self.n_time = self.lambda_NFT.shape
        self.n_mic = self.G_NM.shape[-1]


if __name__ == "__main__":
    import argparse
    import pickle as pic
    import sys, os

    parser = argparse.ArgumentParser()
    parser.add_argument('input_fileName', type= str, help='filename of the multichannel observed signals')
    parser.add_argument(     '--file_id', type= str, default="None", help='file id')
    parser.add_argument(         '--gpu', type= int, default=     0, help='GPU ID')
    parser.add_argument(       '--n_fft', type= int, default=  1024, help='number of frequencies')
    parser.add_argument(    '--n_source', type= int, default=     2, help='number of noise')
    parser.add_argument(     '--n_basis', type= int, default=     8, help='number of basis')
    parser.add_argument(    '--init_SCM', type= str, default="ILRMA", help='unit, obs, ILRMA')
    parser.add_argument( '--n_iteration', type= int, default=   100, help='number of iteration')
    parser.add_argument(   '--save_path', type= str, default="../output/", help='save path of wave file')
    parser.add_argument(    '--dl_ratio', type= float, default=     0, help='diagonal loading ratio')
    parser.add_argument(    '--save_wav', type= str, default=True, help='save wave file')
    args = parser.parse_args()

    if args.gpu < 0:
        import numpy as xp
    else:
        import cupy as xp
        print("Use GPU " + str(args.gpu))
        cuda.get_device_from_id(args.gpu).use()

    wav, fs = sf.read(args.input_fileName)
    wav = wav.T
    M = len(wav)
    for m in range(M):
        tmp = librosa.core.stft(wav[m], n_fft=args.n_fft, hop_length=int(args.n_fft/4))
        if m == 0:
            spec = np.zeros([tmp.shape[0], tmp.shape[1], M], dtype=np.complex)
        spec[:, :, m] = tmp

    separater = FastMNMF2(n_source=args.n_source, n_basis=args.n_basis, xp=xp, init_SCM=args.init_SCM, dl_ratio=args.dl_ratio)
    separater.load_spectrogram(spec)
    separater.file_id = args.file_id
    separater.fs = fs
    separater.solve(n_iteration=args.n_iteration, save_likelihood=False, save_parameter=False, save_wav=args.save_wav, save_path=args.save_path, interval_save_parameter=25)
